from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import BaseOutputParser
from langchain_ollama import ChatOllama
from typing import Dict, Any, ClassVar

router = APIRouter()

# =============================
# 自定义解析器（更健壮：中英文标点、别名、容错）
# =============================
class PoemParser(BaseOutputParser):
    """解析格式：标题:XXX;作者:YYY;内容:ZZZ（兼容中文：标题：XXX；作者：YYY；内容：ZZZ）"""
    EXPECTED_KEYS: ClassVar[list[str]] = ["标题", "作者", "内容"]
    ALIASES: ClassVar[Dict[str, list[str]]] = {
        "标题": ["标题", "题目", "title", "Title"],
        "作者": ["作者", "诗人", "author", "Author"],
        "内容": ["内容", "正文", "全文", "诗文", "content", "Content"],
    }

    def parse(self, text: str) -> Dict[str, Any]:
        norm = (text or "").strip()
        # 统一标点为英文
        norm = norm.replace("；", ";").replace("：", ":")
        parts = [p.strip() for p in norm.split(";") if p.strip()]
        kv: Dict[str, str] = {}
        for part in parts:
            if ":" in part:
                k, v = part.split(":", 1)
                kv[k.strip()] = v.strip()

        # 标准化键名到：标题/作者/内容
        out: Dict[str, Any] = {k: "" for k in self.EXPECTED_KEYS}
        for std_key, aliases in self.ALIASES.items():
            for alias in aliases:
                if alias in kv and kv[alias]:
                    out[std_key] = kv[alias]
                    break
        # 基础校验
        missing = [k for k in self.EXPECTED_KEYS if not out.get(k)]
        if missing:
            raise ValueError(f"缺少必要字段: {','.join(missing)}；原始输出: {text}")
        return out


class QuoteParser(BaseOutputParser):
    """解析格式：名言:XXX;作者:YYY（兼容中文：名言：XXX；作者：YYY）"""
    EXPECTED_KEYS: ClassVar[list[str]] = ["名言", "作者"]
    ALIASES: ClassVar[Dict[str, list[str]]] = {
        "名言": ["名言", "引语", "语录", "quote", "Quote"],
        "作者": ["作者", "人物", "author", "Author"],
    }

    def parse(self, text: str) -> Dict[str, Any]:
        norm = (text or "").strip()
        norm = norm.replace("；", ";").replace("：", ":")
        parts = [p.strip() for p in norm.split(";") if p.strip()]
        kv: Dict[str, str] = {}
        for part in parts:
            if ":" in part:
                k, v = part.split(":", 1)
                kv[k.strip()] = v.strip()

        out: Dict[str, Any] = {k: "" for k in self.EXPECTED_KEYS}
        for std_key, aliases in self.ALIASES.items():
            for alias in aliases:
                if alias in kv and kv[alias]:
                    out[std_key] = kv[alias]
                    break
        missing = [k for k in self.EXPECTED_KEYS if not out.get(k)]
        if missing:
            raise ValueError(f"缺少必要字段: {','.join(missing)}；原始输出: {text}")
        return out


# =============================
# 初始化本地 Ollama 模型
# =============================
llm = ChatOllama(
    model="qwen3:0.6b",
    base_url="http://172.16.21.38:11436"
)

# =============================
# Prompt（ChatPromptTemplate，严格约束格式）
# =============================
poem_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个严格遵循输出格式的助手，只输出我指定的格式，不要添加任何解释或其他文本。"),
    ("human", "请以如下格式输出{poem_name}的诗歌信息：标题:XXX;作者:YYY;内容:ZZZ")
])

quote_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个严格遵循输出格式的助手，只输出我指定的格式，不要添加任何解释或其他文本。"),
    ("human", "请以如下格式输出{person}的一句经典名言：名言:XXX;作者:YYY")
])

# =============================
# 接口 1：解析诗歌信息
# =============================
@router.get("/parse_poem")
async def parse_poem(poem_name: str = Query(..., description="诗歌名称")):
    try:
        messages = poem_prompt.format_messages(poem_name=poem_name)
        resp = llm.invoke(messages)
        text = resp.content if hasattr(resp, "content") else str(resp)
        data = PoemParser().parse(text)
        return JSONResponse(content={"poem": data})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": "解析失败", "details": str(e)})


# =============================
# 接口 2：解析名人名言
# =============================
@router.get("/parse_quote")
async def parse_quote(person: str = Query(..., description="人物姓名")):
    try:
        messages = quote_prompt.format_messages(person=person)
        resp = llm.invoke(messages)
        text = resp.content if hasattr(resp, "content") else str(resp)
        data = QuoteParser().parse(text)
        return JSONResponse(content={"quote": data})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": "解析失败", "details": str(e)})
